package com.cn.lp.ai.factory;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.vectorstore.RedisVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.List;

@Component
public class DocumentService {

    @Value("classpath:meituan-qa.txt")
    private Resource resource;
    @Autowired
    private RedisVectorStore vectorStore;

    /**
     * 向量存储
     * @return
     */
    public List<Document> loadText() {
        //文本读取
        TextReader textReader = new TextReader(resource);
        textReader.getCustomMetadata().put("filename", "meituan-qa.txt");
        List<Document> documents = textReader.get();

        CustomerTextSplitter customerTextSplitter= new CustomerTextSplitter();
        List<Document> list = customerTextSplitter.apply(documents);
        // 把问题存到元数据中
        list.forEach(document -> document.getMetadata().put("question", document.getContent().split("\\n")[0]));
        // 向量存储（文本存储）
        vectorStore.add(list);
        return list;
    }

    /**
     * 向量搜索
     * @param message
     * @return
     */
    public List<Document> search(String message) {
        List<Document> documents = vectorStore.similaritySearch(message);
        return documents;
    }

    /**
     * 元数据搜索
     * @param message
     * @param question
     * @return
     */
    public List<Document> metadataSearch(String message, String question) {
        return vectorStore.similaritySearch(
                SearchRequest
                        .query(message)
//                        .withTopK(5)
                        .withSimilarityThreshold(0.1)
                        .withFilterExpression(String.format("question in ['%s']", question)));
    }

}
